import numpy as np
import torch
import torch.nn as nn


def make_layers(k, input_channel, output_channel1, output_channel2):
    layer = []
    conv1 = nn.Conv2d(in_channels=input_channel, out_channels=output_channel1, kernel_size=1, stride=1)
    conv2 = nn.Conv2d(in_channels=output_channel1, out_channels=output_channel2, kernel_size=k, stride=1, padding=k//2)
    layer.append(conv1)
    layer.append(conv2)
    return nn.Sequential(*layer)

class Inception(nn.Module):
    def __init__(self,input_channel,c1,c2,c3,c4,c5,c6):
        super(Inception,self).__init__()
        self.conv1 = nn.Conv2d(in_channels=input_channel, out_channels=c1, kernel_size=1, stride=1)
        self.conv3 = make_layers(3, input_channel, c2, c3)
        self.conv5 = make_layers(5, input_channel, c4,c5)
        self.pool = nn.MaxPool2d(kernel_size=3,stride=1,padding=1)
        self.conv11 = nn.Conv2d(in_channels=input_channel, out_channels=c6, kernel_size=1, stride=1)
    def forward(self,x):
        x1 = self.conv1(x)
        x3 = self.conv3(x)
        x5 = self.conv5(x)
        x_pool = self.conv11(self.pool(x))
        return torch.concat([x1,x3,x5,x_pool],1)



class GooLeNet(nn.Module):
    def __init__(self):
        super(GooLeNet,self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1,out_channels=64,kernel_size=7,stride=2,padding=3)
        self.pool = nn.MaxPool2d(kernel_size=3,stride=2,padding=1)
        self.conv2 = nn.Conv2d(in_channels=64, out_channels=192, kernel_size=1, stride=1)
        self.conv3 = nn.Conv2d(in_channels=192, out_channels=192, kernel_size=3, stride=1, padding=1)
        self.Inception3a = Inception(192,64,96,128,16,32,32)
        self.Inception3b = Inception(256, 128, 128, 192, 32, 96, 64)
        self.Inception4a = Inception(480, 192, 96, 208, 16, 48, 64)
        self.Inception4b = Inception(512, 160, 112, 224, 24, 64, 64)
        self.Inception4c = Inception(512, 128, 128, 256, 24, 64, 64)
        self.Inception4d = Inception(512, 112, 144, 288, 32, 64, 64)
        self.Inception4e = Inception(528, 256, 160, 320, 32, 128, 128)
        self.Inception5a = Inception(832, 256, 160, 320, 32, 128, 128)
        self.Inception5b = Inception(832, 384, 192, 384, 48, 128, 128)
        self.avgpool = nn.AvgPool2d(kernel_size=7)
        self.dropout = nn.Dropout(0.4)
        self.linear = nn.Linear(in_features=1024,out_features=2)
        self.softmax = nn.Softmax()
    def forward(self,x):
        x = self.conv1(x)
        x = self.pool(x)
        x = self.conv2(x)
        x = self.pool(self.conv3(x))
        x = self.Inception3a(x)
        x = self.Inception3b(x)
        x = self.pool(x)
        x = self.Inception4a(x)
        x = self.Inception4b(x)
        x = self.Inception4c(x)
        x = self.Inception4d(x)
        x = self.Inception4e(x)
        x = self.pool(x)
        x = self.Inception5a(x)
        x = self.Inception5b(x)
        x = self.avgpool(x)
        x = x.view(-1,x.size(1))
        x = self.dropout(x)
        x = self.softmax(self.linear(x))


        return x




#input = torch.randn(1,3,224,224)
#model = GooLeNet()
#output = model(input)
#print(output.shape)
#total = sum([param.nelement() for param in model.parameters()])

#print("Number of parameter: %.2fM" % (total / 1e6))
